'''
https://leetcode.cn/problems/accounts-merge/?envType=study-plan-v2&envId=graph-theory
'''
from collections import defaultdict
from typing import List

class UF:
    def __init__(self, n):
        self.father = list(range(n))
    def find(self, x):
        if x != self.father[x]:
            self.father[x] = self.find(self.father[x])
        return self.father[x]
    def union(self, x, y):
        fx, fy = self.find(x), self.find(y)
        if fx != fy:
            self.father[fy] = fx

class Solution:
    def accountsMerge(self, accounts: List[List[str]]) -> List[List[str]]:
        id = 0
        email_id = {}
        id_name = {}
        for account in accounts:
            name = account[0]
            for email in account[1:]:
                email_id[email] = id
                id_name[id] = name
                id += 1

        uf = UF(id)
        for account in accounts:
            first_email_id = email_id[account[1]]
            for email in account[2:]:
                uf.union(first_email_id, email_id[email])
        id_emails = defaultdict(list)
        for email,id in email_id.items():
            id_emails[uf.find(id)].append(email)
        res = []
        for id, emails in id_emails.items():
            res.append([id_name[id]] + sorted(emails))
        return res

accounts= [["John","johnsmith@mail.com","john_newyork@mail.com"],["John","johnsmith@mail.com","john00@mail.com"],["Mary","mary@mail.com"],["John","johnnybravo@mail.com"]]
print(Solution().accountsMerge(accounts))